using UnityEngine;
using System.Collections;
using AForge.Neuro;
using AForge.Neuro.Learning;

public class TestScript : MonoBehaviour {
	
	public Transform ball1;
	public Transform ball2;
	public Transform mouse;
	
	public const float MouseMaxSpeed = 8f;
	public const float BallMaxSpeed = 8f;
	const float EastWestWallDist = 5.0f;
	const float NorthSouthWallDist = 5.0f;
	
	private BackPropagationLearning teacher;
	private ActivationNetwork nn;
	const float SigmoidAlpha = 1.0f;
	const int DimInput = 10;
	const int FirstLayerNeurons = 7;
	const int SecondLayerNeurons = 8;
	const int DimOutput = 2;
	const float LearningRate = 0.01f;
	const float Momentum = 0.9f;
	//const float ErrorTolerance = 0.1f;
	const int NumEpochs = 10000;
	
	const int TrainingSetSize = 5000;//240;
	private int TrainingSetIndex = 0;
	private float TrainingSamplePeriod = 0.5f; // measured in seconds
	const float TrainingStartTime = 3f; // measured in seconds
	private double[][] TrainingInput = new double[TrainingSetSize][];
	private double[][] TrainingOutput = new double[TrainingSetSize][];
	private float LastTime;
	
	private bool IsTraining = true;
	private bool IsManualTraining = false;
	
	
	// Use this for initialization
	void Start () {
		LastTime = Time.time;
		
		ball1.rigidbody.AddForce (new Vector3(8,0,8), ForceMode.Impulse);
		ball2.rigidbody.AddForce (new Vector3(8,0,-8), ForceMode.Impulse);
		ball2.rigidbody.AddTorque (new Vector3(0.5f,0,0.5f), ForceMode.Impulse);
		
		nn = new ActivationNetwork(new BipolarSigmoidFunction(SigmoidAlpha),
			DimInput, FirstLayerNeurons, SecondLayerNeurons, DimOutput);
		
		teacher = new BackPropagationLearning(nn);
		teacher.LearningRate = LearningRate;
		teacher.Momentum = Momentum;
	}
	
	// Update is called once per frame
	void Update () {
		
		if (Time.time < TrainingStartTime)
		{
			mouse.GetComponent<MouseScript>().DesiredPosition = RandomDestination(mouse.position, 2f);
			
			if (Input.GetButton("Fire1"))
			{
				//IsManualTraining = true;
				
				//TrainingSamplePeriod = 0.25f;
			}
			return;
		}
		
		if (IsTraining) // Training?
		{
			MouseScript m = mouse.GetComponent<MouseScript>();
			Vector3 TrainingVelocity;
			if (IsManualTraining)
			{
				TrainingVelocity = GetInput();
				//m.DesiredVelocity = GetInput();
				print(Input.mousePosition);
			}
			else
			{
				TrainingVelocity = PotentialFieldForce();
			}
			
			TrainingVelocity = new Vector3(TrainingVelocity.x, 0, TrainingVelocity.z);
			
			if ((m.transform.position - m.DesiredPosition).sqrMagnitude < 0.5f)
			{
				m.DesiredPosition = RandomDestination(m.transform.position, 3f);
			}
			m.DesiredVelocity = (m.DesiredPosition - m.transform.position).normalized * MouseMaxSpeed;
			
			
			if (TrainingSetIndex >= TrainingSetSize)
			{
				print ("Data acquired!");
				IsTraining = false;
				double error = double.MaxValue;
				error = 0;
				for (int i = 0; i < NumEpochs; i++)
				{
					error = teacher.RunEpoch(TrainingInput, TrainingOutput) / NumEpochs;
				}
				print (error);
			}
			else if (Time.time > LastTime + TrainingSamplePeriod)
			{
				// make training set
				TrainingInput[TrainingSetIndex] = new double[DimInput];
				TrainingOutput[TrainingSetIndex] = new double[DimOutput];
				
				Vector3 ball1Vel = ball1.rigidbody.velocity.normalized;
				Vector3 ball2Vel = ball2.rigidbody.velocity.normalized;
				
				TrainingInput[TrainingSetIndex][0] = mouse.position.x / EastWestWallDist;
				TrainingInput[TrainingSetIndex][1] = mouse.position.z / NorthSouthWallDist;
				TrainingInput[TrainingSetIndex][2] = ball1.position.x / EastWestWallDist;
				TrainingInput[TrainingSetIndex][3] = ball1.position.z / NorthSouthWallDist;
				TrainingInput[TrainingSetIndex][4] = ball1Vel.x;
				TrainingInput[TrainingSetIndex][5] = ball1Vel.z;
				TrainingInput[TrainingSetIndex][6] = ball2.position.x / EastWestWallDist;
				TrainingInput[TrainingSetIndex][7] = ball2.position.z / NorthSouthWallDist;
				TrainingInput[TrainingSetIndex][8] = ball2Vel.x;
				TrainingInput[TrainingSetIndex][9] = ball2Vel.z;
				//TrainingInput[TrainingSetIndex][6] = Mathf.Atan2(ballVel.z, ballVel.x) / Mathf.PI * 2.0f;
				
				// normalize input
				//NormalizeInput(TrainingInput[TrainingSetIndex]);
				
				LastTime = Time.time;
				
				if (!(TrainingVelocity.sqrMagnitude < 0.1f && IsManualTraining))
				{
					Vector3 desiredVelocity = TrainingVelocity;
					desiredVelocity.Normalize();
					TrainingOutput[TrainingSetIndex][0] = desiredVelocity.x;
					TrainingOutput[TrainingSetIndex][1] = desiredVelocity.z;
				
					print(TrainingSetIndex++);
				}
			}
		}
		else // Playing
		{
			double[] testingInput = new double[DimInput];
			
			Vector3 ball1Vel = ball1.rigidbody.velocity.normalized;
			Vector3 ball2Vel = ball2.rigidbody.velocity.normalized;
			
			testingInput[0] = mouse.position.x / EastWestWallDist;
			testingInput[1] = mouse.position.z / NorthSouthWallDist;
			testingInput[2] = ball1.position.x / EastWestWallDist;
			testingInput[3] = ball1.position.z / NorthSouthWallDist;
			testingInput[4] = ball1Vel.x;
			testingInput[5] = ball1Vel.z;
			testingInput[6] = ball2.position.x / EastWestWallDist;
			testingInput[7] = ball2.position.z / NorthSouthWallDist;
			testingInput[8] = ball2Vel.x;
			testingInput[9] = ball2Vel.z;
			//testingInput[6] = Mathf.Atan2(ballVel.z, ballVel.x) / Mathf.PI * 2.0f;
			
			// normalize input
			//NormalizeInput(testingInput);
			
			double[] testingOutput = nn.Compute(testingInput);
			mouse.GetComponent<MouseScript>().DesiredVelocity = new Vector3(
				MouseMaxSpeed * (float)testingOutput[0], 0, MouseMaxSpeed * (float)testingOutput[1]);
			
		}
	}
	
	void NormalizeInput(double[] input)
	{
		double max = 0;
		for (int i = 0; i < DimInput; i++)
		{
			if (max < input[i])
				max = input[i];
		}
		if (max == 0) return;
		for (int i = 0; i < DimInput; i++)
		{
			input[i] /= max;
		}
	}
	
	Vector3 GetInput()
	{
		float v = Input.GetAxis("Mouse Y");
		float h = Input.GetAxis("Mouse X");
		if (v == h && h == 0)
		{
			v = Input.GetAxis("Vertical");
			h = Input.GetAxis("Horizontal");			
		}
		
		Vector3 u = new Vector3(h, 0, v);
		u *= MouseMaxSpeed;
		return u;
	}
	
	Vector3 RandomDestination(Vector3 origin, float range)
	{
		Vector3 position;
		do
		{
			position = origin + Random.insideUnitSphere * range;
		}
		while (position.x > EastWestWallDist - 0.5f || position.x < -EastWestWallDist + 0.5f ||
			   position.z > NorthSouthWallDist - 0.5f || position.z < -NorthSouthWallDist + 0.5f);
		
		return new Vector3(position.x, origin.y, position.z);
	}
	
	Vector3 PotentialFieldForce()
	{
		const float ballK = 4.5f;
		const float ballVK = 0.2f;
		const float wallK = 0.2f;
		const float attractK = 0.5f;
		
		Vector3 dir1 = mouse.position - ball1.position;
		dir1 = new Vector3(dir1.x, 0, dir1.z);
		Vector3 ball1Force = dir1.normalized * ballK / (dir1.magnitude);//Vector3.Dot(dir1, dir1));
		Vector3 dir1V = mouse.position - ball1.position - ball1.rigidbody.velocity * ballVK;
		dir1V = new Vector3(dir1V.x, 0, dir1V.z);
		Vector3 ball1VForce = dir1V.normalized * ballK / (Vector3.Dot(dir1V, dir1V));
		
		Vector3 dir2 = mouse.position - ball2.position;
		dir2 = new Vector3(dir2.x, 0, dir2.z);
		Vector3 ball2Force = dir2.normalized * ballK / dir2.magnitude;//(Vector3.Dot(dir2, dir2));
		Vector3 dir2V = mouse.position - ball2.position - ball2.rigidbody.velocity * ballVK;
		dir2V = new Vector3(dir2V.x, 0, dir2V.z);
		Vector3 ball2VForce = dir2V.normalized * ballK / (Vector3.Dot(dir2V, dir2V));
		
		Vector3 ballAttraction = -attractK * dir2;
		
		/*
		Vector3 eastForce = new Vector3(wallK / (mouse.position.x - EastWestWallDist) / (mouse.position.x - EastWestWallDist), 0, 0);
		Vector3 westForce = new Vector3(wallK / (mouse.position.x + EastWestWallDist) / (mouse.position.x + EastWestWallDist), 0, 0);
		Vector3 northForce = new Vector3(0, 0, wallK / (mouse.position.z - NorthSouthWallDist) / (mouse.position.z - NorthSouthWallDist));
		Vector3 southForce = new Vector3(0, 0, wallK / (mouse.position.z + NorthSouthWallDist) / (mouse.position.z + NorthSouthWallDist));
		
		Vector3 nwDir = mouse.position - new Vector3(-EastWestWallDist, 0, NorthSouthWallDist);
		Vector3 nwCornerForce = wallK * 2 / nwDir.sqrMagnitude * nwDir.normalized;
		Vector3 neDir = mouse.position - new Vector3(EastWestWallDist, 0, NorthSouthWallDist);
		Vector3 neCornerForce = wallK * 2 / neDir.sqrMagnitude * neDir.normalized;
		Vector3 swDir = mouse.position - new Vector3(-EastWestWallDist, 0, -NorthSouthWallDist);
		Vector3 swCornerForce = wallK * 2 / swDir.sqrMagnitude * swDir.normalized;
		Vector3 seDir = mouse.position - new Vector3(EastWestWallDist, 0, -NorthSouthWallDist);
		Vector3 seCornerForce = wallK * 2 / seDir.sqrMagnitude * seDir.normalized;
		*/
		Vector3 wallForce = -wallK * mouse.position.sqrMagnitude * mouse.position.normalized;
		
		
		Vector3 total = ball1Force + ball1VForce + ball2Force + ball2VForce
						+ wallForce;//- eastForce + westForce - northForce + southForce;
		
		//float maxWallForce = Mathf.Max(new float[]{eastForce.sqrMagnitude, 
		//	westForce.sqrMagnitude, northForce.sqrMagnitude, southForce.sqrMagnitude});
		
		TrainingSamplePeriod = Mathf.Min(2f / (ball1Force.sqrMagnitude + ball2Force.sqrMagnitude + wallForce.sqrMagnitude), 2f);
		
		//print ("maxwallforce: " + wallForce.sqrMagnitude + " ball1force: " + ball1Force.sqrMagnitude + "TrainingSamplePeriod" + TrainingSamplePeriod);
		
		if (total.sqrMagnitude > MouseMaxSpeed * MouseMaxSpeed)
			total = total.normalized * MouseMaxSpeed;
			
		return total;
	}
}
